import time
import socket
import struct


class IpRegion:
    Version = 2
    IndexPolicy = 1
    HeaderSize = 256
    VectorIndexRows = 256
    VectorIndexCols = 256
    VectorIndexSize = 8
    SegmentIndexSize = 14
    VectorIndexesSize = VectorIndexRows * VectorIndexCols * VectorIndexSize

    IP_INFO_TYPES = 4  # IP_START | IP_END | 国家 | 城市 | 运营商

    @staticmethod
    def is_ip(ip: str):
        ip_segments = ip.strip().split(".")
        if len(ip_segments) != 4:
            return False
        for ip_segment in ip_segments:
            if not ip_segment.isdigit():
                return False
            if len(ip_segment) > 3:
                return False
            if int(ip_segment) > 255:
                return False
        return True

    @staticmethod
    def ip_to_bytes(ip: str):
        """ip转int"""
        return socket.inet_aton(ip)

    @staticmethod
    def bytes2int(byte):
        """四字节"""
        return struct.unpack('I', byte)[0]

    @staticmethod
    def bytes2long(byte):
        """四字节"""
        return struct.unpack('!L', byte)[0]

    @staticmethod
    def ip_segment(start_ip, end_ip):
        """
        ip 大段拆分成小段
        :param start_ip: 116.31.76.0
        :param end_ip:  117.21.79.49
        :return:
        (116.31.76.0 , 116.31.255.255)
        (116.32.0.0  , 116.32.255.255)
         ...         , ...
        (116.255.0.0 , 116.255.255.255 )
        (117.0.0.0   , 117.0.255.255  )
        (117.1.0.0   , 117.1.255.255  )
         ...         , ...
        (117.21.0.0  , 117.21.79.49 )
        """
        t_list_1 = []
        s_byte_1, e_byte_1 = (start_ip >> 24) & 0xFF, (end_ip >> 24) & 0xFF
        n_sip = start_ip
        for i in range(s_byte_1, e_byte_1 + 1):
            sip = (i << 24) | (n_sip & 0xFFFFFF)
            eip = (i << 24) | 0xFFFFFF
            if eip < end_ip:
                n_sip = (i + 1) << 24
            else:
                eip = end_ip
            # Append the new segment (maybe)
            t_list_1.append((sip, eip))

        # 2. Split the segments with the second byte
        t_list_2 = []
        for s in t_list_1:
            base = s[0] & 0xFF000000
            n_sip = s[0]
            s_byte_2, e_byte_2 = (s[0] >> 16) & 0xFF, (s[1] >> 16) & 0xFF
            for i in range(s_byte_2, e_byte_2 + 1):
                sip = base | (i << 16) | (n_sip & 0xFFFF)
                eip = base | (i << 16) | 0xFFFF
                if eip < end_ip:
                    n_sip = 0
                else:
                    eip = end_ip
                t_list_2.append((sip, eip))
        return t_list_2


class Ip2DB(IpRegion):
    """文本写入到数据库"""
    def __init__(self, txt_path, db_path):
        self.region_dict = dict()
        self.start_ptr = -1
        self.end_ptr = -1
        self.vector_index = {(r, c): {'start_ptr': 0, 'end_ptr': 0}
                             for r in range(self.VectorIndexRows) for c in range(self.VectorIndexCols)}
        self.src_handler = open(txt_path, 'r')
        self.db_handler = open(db_path, 'wb')

    def load_segments(self):
        last = None
        segments = []
        lines = self.src_handler.read().splitlines()
        for idx, line in enumerate(lines):
            frags = line.split('|', maxsplit=2)
            if len(frags) != 3:
                raise Exception(f'第{idx}行: 无效的数据格式')

            if not self.is_ip(frags[0]):
                raise Exception(f'第{idx}行: 无效的开始IP')

            if not self.is_ip(frags[1]):
                raise Exception(f'第{idx}行: 无效的结束IP')

            start_ip, end_ip = self.bytes2long(self.ip_to_bytes(frags[0])), self.bytes2long(self.ip_to_bytes(frags[1]))
            if start_ip > end_ip:
                raise Exception(f'第{idx}行: 无效的开始IP与结束IP')

            if len(frags[2]) < 1:
                raise Exception(f'第{idx}行: 无效的地址数据信息')

            if len(bytes(frags[2], encoding='utf-8')) > 0xFFFF:
                raise Exception(f'第{idx}行: 地址数据信息量过大')

            if last is not None:
                assert last + 1 == end_ip, f'第{idx}行: IP地址段不连续'
                last = end_ip

            segments.append(frags)
        return segments

    def write_header(self):
        header = bytearray([0] * 256)
        header[0:2] = self.Version.to_bytes(2, byteorder="little")
        header[2:4] = self.IndexPolicy.to_bytes(2, byteorder="little")
        header[4:8] = int(time.time()).to_bytes(4, byteorder="little")
        # 4. Index block start ptr
        header[8:12] = int(0).to_bytes(4, byteorder="little")
        # 5. Index block end ptr
        header[12:16] = int(0).to_bytes(4, byteorder="little")
        self.db_handler.write(header)

    def _get_vector_index(self, ip, ptr):
        row, col = (ip >> 24) & 0xFF, (ip >> 16) & 0xFF
        vector_index = self.vector_index[(row, col)]
        if vector_index['start_ptr'] == 0:
            vector_index['start_ptr'] = ptr
        vector_index['end_ptr'] = ptr + self.SegmentIndexSize

    def write_vector(self, segments):
        for segment in segments:
            region_byte = bytes(segment[2], encoding="utf-8")
            region_len = len(region_byte)

            start_ip = self.bytes2long(self.ip_to_bytes(segment[0]))
            end_ip = self.bytes2long(self.ip_to_bytes(segment[1]))
            _segments = self.ip_segment(start_ip, end_ip)

            for _segment in _segments:  # 二分索引地址
                position = self.db_handler.seek(0, 1)
                index = struct.pack("<IIHI", _segment[0], _segment[1], region_len, self.region_dict[segment[2]])
                self.db_handler.write(index)

                self._get_vector_index(_segment[0], position)
                if self.start_ptr == -1:
                    self.start_ptr = position
                self.end_ptr = position

        self.db_handler.seek(self.HeaderSize, 0)
        for i in range(0, self.VectorIndexRows):
            for j in range(0, self.VectorIndexRows):
                vector_index = self.vector_index[(i, j)]
                index = struct.pack("<II", vector_index['start_ptr'], vector_index['end_ptr'])
                self.db_handler.write(index)

    def write_region(self, segments):
        self.db_handler.seek(self.HeaderSize + self.VectorIndexesSize, 0)
        for segment in segments:
            if segment[2] in self.region_dict:
                continue
            region_byte = bytes(segment[2], encoding="utf-8")
            position = self.db_handler.seek(0, 1)
            self.db_handler.write(region_byte)
            self.region_dict[segment[2]] = position

    def write_index(self):
        buff = struct.pack("<II", self.start_ptr, self.end_ptr)
        self.db_handler.seek(8, 0)
        self.db_handler.write(buff)

    def close(self):
        self.src_handler.close()
        self.db_handler.close()

    def translate(self):
        segments = self.load_segments()

        self.write_header()
        self.write_region(segments)
        self.write_vector(segments)
        self.write_index()

        self.close()


class Searcher(IpRegion):
    def __init__(self, db_path):
        self.db_path = db_path
        with open(self.db_path, 'rb') as f:
            self.content = f.read()

    @staticmethod
    def bytes16_to_long(byte):
        """十六字节"""
        return struct.unpack('IL', byte)[0]

    @staticmethod
    def get_position(b, offset):
        return (b[offset] & 0x000000FF) | (b[offset+1] << 8)

    def ip_index(self, ip: int):
        il0 = int((ip >> 24) & 0xFF)
        il1 = int((ip >> 16) & 0xFF)
        idx = il0 * self.VectorIndexCols * self.VectorIndexSize + il1 * self.VectorIndexSize
        return idx

    def get_buffer(self, offset, size):
        buffer = self.content[offset:offset + size]
        return buffer

    def search(self, ip):
        byte_ip = self.ip_to_bytes(ip)
        ip_long = self.bytes2long(byte_ip)

        idx = self.ip_index(ip_long)
        _idx = self.HeaderSize + idx
        s_pos = _idx + 4
        e_pos = s_pos + 4

        s_ptr = self.bytes2int(self.content[_idx:s_pos])
        e_ptr = self.bytes2int(self.content[s_pos:e_pos])

        dataLen = dataPtr = 1
        l = 0
        h = int((e_ptr - s_ptr) / self.SegmentIndexSize)
        while l <= h:
            m = int((l + h) >> 1)
            p = int(s_ptr + m * self.SegmentIndexSize)

            buffer_sip = self.get_buffer(p, self.SegmentIndexSize)
            sip = self.bytes2int(buffer_sip[0:4])
            if ip_long < sip:
                h = m - 1
            else:
                eip = self.bytes2int(buffer_sip[4:8])
                if ip_long > eip:
                    l = m + 1
                else:
                    dataLen = self.get_position(buffer_sip, 8)
                    dataPtr = self.bytes2int(buffer_sip[10:14])
                    break

        if dataPtr < 0:
            return ""

        buffer_string = self.get_buffer(dataPtr, dataLen)
        return_string = buffer_string.decode("utf-8")
        return return_string


if __name__ == '__main__':

    # ip2db = Ip2DB('./data/ip.merge.txt', './data/ip.xdb')
    # ip2db.translate()

    ip_1 = '118.116.111.23'
    searcher = Searcher(db_path='./data/ip2region.xdb')
    region = searcher.search(ip_1)
    print(region)
